import sys
import socket
import getopt
import threading
import subprocess

# define some global variable
listen = False
command = False
upload = False
execute = ""
target = ""
upload_destination = ""
port = 0


def usage():
    print "BHP Net Tool"
    print
    print "Usage: bhpnet.py -t target_host -p port"
    print "-l --listen                  - listen on [host]:[port] for" \
          "                               incoming connections"
    print "-e --execute=file_to_run     - execute the given file upon" \
          "                               receiving a connection"
    print "-c --command                 - initialize a command shell"
    print "-u --upload=destination      - upon receiving connection upload a " \
          "                               file and write to [destination]"
    print
    print

    print "Examples: "
    print "bhpnet.py -t 192.168.0.1 -p 5555 -l -c"
    print "bhpnet.py -t 192.168.0.1 -p 5555 -l -u=C:\\target.exe"
    print "bhpnet.py -t 192.168.0.1 -p 5555 -l -e=\"cat /etc/passwd\""
    print "echo 'ABCDEFGHI' | ./bhpnet.py -t 192.168.11.12 -p 135"
    sys.exit(0)


def client_sender(buffer):
    client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

    try:
        # connect to target
        client.connect((target, port))

        if len(buffer):
            client.send(buffer)

        while True:
            # wait for data feedback
            recv_len = 1
            response = ""

            while recv_len:
                data = client.recv(4096)
                recv_len = len(data)
                response += data

                if recv_len < 4096:
                    break

            print response,

            # wait more feedback
            buffer = raw_input("")
            buffer += "\n"

            #send
            client.send(buffer)
    except:
        print "[*] Exception! Exiting."
        # close connect
        client.close()


def server_loop():
    global target

    # if no define target, listening all port
    if not len(target):
        target = "0.0.0.0"

    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.bind((target, port))

    server.listen(5)
    print "server start listening at %s:%d" % (target, port)

    while True:
        client_socket, addr = server.accept()

        # create a thread to handle new client
        client_thread = threading.Thread(target=client_handler,
                                         args=(client_socket,))
        client_thread.start()


def run_command(command):
    # newline
    command = command.rstrip()

    # run command and print feedback
    try:
        output = subprocess.check_output(command, stderr=subprocess.STDOUT, shell=True)
    except:
        output = "Failed to execute command.\r\n"

    #send output
    return output


def client_handler(client_socket):
    global upload
    global execute
    global command

    print "==========================start client_handler=========================="
    print "upload: %s, execute: %s, command: %s" % (upload, execute, command)
    # check upload
    if len(upload_destination):
        print "in upload handle"
        # read all data and write target
        file_buffer = ""

        # read data still no data
        while True:
            data = client_socket.recv(1024)

            if not data:
                break
            else:
                file_buffer += data

        # write the data that we received
        try:
            file_descriptor = open(upload_destination, "wb")
            file_descriptor.write(file_buffer)
            file_descriptor.close()

            #confirm file write right
            client_socket.send("Successfully saved file to %s\r\n" % upload_destination)
        except:
            client_socket.send("Failed to save file to %s\r\n" % upload_destination)

    # check command execute
    if len(execute):
        print "in execute handle"
        # run command
        output = run_command(excute)
        client_socket.send(output)

    # if we need a command line shell, go to another while
    if command:
        print "in command handle"
        while True:
            # create a new window
            client_socket.send("<BHP:#> ")
            # receive file still enter key
            cmd_buffer = ""
            while "\n" not in cmd_buffer:
                cmd_buffer += client_socket.recv(1024)
            # return output
            response = run_command(cmd_buffer)

            # send feedback
            client_socket.send(response)


def main():
    global listen
    global port
    global execute
    global command
    global upload_destination
    global target

    if not len(sys.argv[1:]):
        usage()

    # read command line argv
    try:
        opts, args = getopt.getopt(sys.argv[1:], "hle:t:p:cu:",
                                   ["help", "listen", "execute", "target", "port", "command", "upload"])
    except getopt.GetoptError as err:
        print str(err)
        usage()

    for o, a in opts:
        if o in ("-h", "--help"):
            usage()
        elif o in ("-l", "--listen"):
            print "-l change listen to True"
            listen = True
        elif o in ("-e", "--execute"):
            print "-e change execute to %s" % a
            execute = a
        elif o in ("-c", "--commandshell"):
            print "-c change command to True"
            command = True
        elif o in ("-u", "--upload"):
            print "-u change upload to %s" % a
            upload_destination = a
        elif o in ("-t", "--target"):
            print "-t change target to %s" % a
            target = a
        elif o in ("-p", "--port"):
            print "-p change port to %s" % a
            port = int(a)
        else:
            assert False, "Unhandled Option"

    # choose listening or send data from inputstream
    if not listen and len(target) and port > 0:
        print "=======================command mode======================================"
        # read data from RAM by commandline
        # if you do this , there will be block , you should send CTRL -D when you don't want to send data again
        buffer_test = sys.stdin.read()

        print "\nreceived: %s" % buffer_test
        # send data
        client_sender(buffer_test)

    if listen:
        print "=======================listening mode======================================"
        server_loop()

main()
